#need RBTreeBASE.py
from RBTreeBASE import node,BST
from seg import seg
class RBTree(BST):
    def __init__(self):
        self.NIL=node(key='NIL')
        self.root=self.NIL
    def interval_search(self, i):
        x = self.root
        while x is not self.NIL and not x.key.cross(i):
            if x.left is not self.NIL and x.left.m >= i.lo:
                x = x.left
            else:
                x = x.right
        return x
    def all_interval_search(self, i):
        return self.all_interval_search_from(self.root, i)
    def all_interval_search_from(self, x, i):
        if x is self.NIL: 
            return
        if x.left is not self.NIL and x.left.m >= i.lo:
            self.all_interval_search_from(x.left, i)
        if x.key.cross(i): 
            print x
        if x.right is not self.NIL and x.right.m >= i.lo:
            self.all_interval_search_from(x.right, i)
    def inorder(self):
        BST.inorder(self,NIL=self.NIL)
    def rb_search(self,k):
        return self.search(k,NIL=self.NIL)
    def left_rotate(self,x):
        y=x.right
        x.right=y.left
        x.right.p=x
        y.p=x.p
        if y.p is self.NIL:
            self.root=y
        else:
            if x is x.p.left:
                x.p.left=y
            else:
                x.p.right=y
        y.left=x
        y.left.p=y
        y.m = x.m
        x.m = max(x.left.m, x.right.m, x.key.hi)
    def right_rotate(self,x):
        y=x.left
        x.left=y.right
        x.left.p=x
        y.p=x.p
        if y.p is self.NIL:
            self.root=y
        else:
            if x is x.p.left:
                x.p.left=y
            else:
                x.p.right=y
        y.right=x
        y.right.p=y
        y.m = x.m
        x.m = max(x.left.m, x.right.m, x.key.hi)
    def rb_insert(self,k):
        y=self.NIL
        x=self.root
        while x is not self.NIL:
            x.m = k.hi if k.hi > x.m else x.m
            y=x
            if k<y.key:
                x=x.left
            else:
                x=x.right
        z=node(key=k,p=y,left=self.NIL,right=self.NIL,color='red',m=k.hi)
        if y is self.NIL:
            self.root=z
        else:
            if k<y.key:
                y.left=z
            else:
                y.right=z
        self.rb_insert_fixup(z)
    def rb_insert_fixup(self,z):
        while z.p.color=='red': #if z.p is root, its color is black
            if z.p is z.p.p.left:
                y=z.p.p.right
                if y.color=='red':
                    z.p.color='black'
                    y.color='black'
                    z.p.p.color='red'
                    z=z.p.p
                else:
                    if z is z.p.right:
                        z=z.p
                        self.left_rotate(z)
                    z.p.color='black'
                    z.p.p.color='red'
                    self.right_rotate(z.p.p)
            else: #z.p is z.p.p.right
                y=z.p.p.left
                if y.color=='red':
                    z.p.color='black'
                    y.color='black'
                    z.p.p.color='red'
                    z=z.p.p
                else:
                    if z is z.p.left:
                        z=z.p
                        self.right_rotate(z)
                    z.p.color='black'
                    z.p.p.color='red'
                    self.left_rotate(z.p.p)
        self.root.color='black'
    
if __name__=='__main__':
    tree=RBTree();
    A=[(0,3),(5,8),(6,10),(8,9),(15,23),
       (16,21),(17,19),(19,20),(25,30),(26,26)]
    for i in A:
        tree.rb_insert(seg(*i))
    print 'tree inorder:'
    tree.inorder()
    print
    z=tree.interval_search(seg(16,25))
    print 'seg corss with [16,25]:', z
    print 'segs corss with [16,25]:' 
    z=tree.all_interval_search(seg(16,25))





